package com.example.logTracer.config;

import org.slf4j.MDC;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;
import java.util.concurrent.ThreadPoolExecutor;

/**
 * @Describe : 线程池配置
 * @Author : zyx
 * @Date : 2023/11/17 16:03
 * @Version : 1.1
 **/
@Configuration
@EnableAsync
public class ThreadPoolTaskConfig {

    @Value("${threadPool.corePoolSize}")
    private int corePoolSize;

    @Value("${threadPool.maxPoolSize}")
    private int maxPoolSize;

    @Value("${threadPool.keepAliveSeconds}")
    private int keepAliveSeconds;

    @Bean("poolTask")
    public ThreadPoolTaskExecutor threadPoolTaskExecutor(){
        ThreadPoolTaskExecutor threadPoolTaskExecutor = new HydraThreadPoolTaskExecutor();

        threadPoolTaskExecutor.setCorePoolSize(corePoolSize);

        threadPoolTaskExecutor.setMaxPoolSize(maxPoolSize);

        threadPoolTaskExecutor.setQueueCapacity(1024);

        threadPoolTaskExecutor.setKeepAliveSeconds(keepAliveSeconds);

        threadPoolTaskExecutor.setThreadNamePrefix("async-task-");
        threadPoolTaskExecutor.setWaitForTasksToCompleteOnShutdown(true);

        threadPoolTaskExecutor.setRejectedExecutionHandler(new ThreadPoolExecutor.AbortPolicy());
        threadPoolTaskExecutor.initialize();
        return threadPoolTaskExecutor;
    }

    static class HydraThreadPoolTaskExecutor extends ThreadPoolTaskExecutor{
        private static final long serialVersionUID = 1L;
        private boolean useFixedContext = false;
        private Map<String,String> fixedContext;

        public HydraThreadPoolTaskExecutor(){
            super();
        }

        public HydraThreadPoolTaskExecutor(boolean useFixedContext, Map<String, String> fixedContext) {
            this.useFixedContext = useFixedContext;
            this.fixedContext = fixedContext;
        }

        private Map<String,String> getContextForTask(){
            return useFixedContext ? fixedContext:MDC.getCopyOfContextMap();
        }

        @Override
        public void execute(Runnable task) {
            super.execute(wrapExecute(task,getContextForTask()));
        }

        @Override
        public <T> Future<T> submit(Callable<T> task) {
            return super.submit(wrapSubmit(task,getContextForTask()));
        }

        private <T> Callable<T> wrapSubmit(Callable<T> task ,final Map<String,String> context){
            return ()->{
                Map<String,String> previous = MDC.getCopyOfContextMap();
                if (context == null){
                    MDC.clear();
                }else {
                    MDC.setContextMap(context);
                }

                try {
                    return task.call();
                }finally {
                    if (previous == null){
                        MDC.clear();
                    }else {
                        MDC.setContextMap(previous);
                    }
                }
            };
        }

        private Runnable wrapExecute(final Runnable runnable,final Map<String,String> context){
            return ()->{
                Map<String,String> previous = MDC.getCopyOfContextMap();
                if (context == null){
                    MDC.clear();
                }else {
                    MDC.setContextMap(context);
                }

                try {
                    runnable.run();
                }finally {
                    if (previous == null){
                        MDC.clear();
                    }else {
                        MDC.setContextMap(previous);
                    }
                }
            };
        }
    }

}
